"""
增强版音频-视觉事件数据集模块
该模块定义了支持噪声数据和多标签的AVEDataset类
用于加载和处理包含背景噪声的音频-视觉事件数据集
"""

# 导入操作系统接口模块，提供文件路径操作功能
import os
# 导入HDF5文件操作库，用于处理大型科学数据集
import h5py
# 导入PyTorch深度学习框架
import torch
# 从PyTorch数据工具包中导入数据集和数据加载器相关类
from torch.utils.data import Dataset, DataLoader

# 定义增强版音频-视觉事件数据集类，继承自PyTorch的Dataset类
class AVEDataset(Dataset):
    """
    增强版音频-视觉事件数据集类
    支持干净数据和噪声数据，包含视频级标签和原始标签
    适用于多实例学习任务
    """
    
    def __init__(self, data_root, split='train'):
        """
        初始化数据集
        
        参数:
        data_root: 数据根目录路径，包含所有特征文件和标签文件
        split: 数据集分割类型，默认为'train'，可选'train'、'test'等
        """
        
        # 调用父类Dataset的初始化方法
        super(AVEDataset, self).__init__()
        # 设置数据集分割类型（训练集或测试集）
        self.split = split
        # 构建干净视觉特征文件的完整路径
        self.visual_feature_path = os.path.join(data_root, 'visual_feature.h5')
        # 构建干净音频特征文件的完整路径
        self.audio_feature_path = os.path.join(data_root, 'audio_feature.h5')
        # 构建噪声视觉特征文件的完整路径（仅包含背景）
        self.noisy_visual_feature_path = os.path.join(data_root, 'visual_feature_noisy.h5')
        # 构建噪声音频特征文件的完整路径（仅包含背景）
        self.noisy_audio_feature_path = os.path.join(data_root, 'audio_feature_noisy.h5')
        
        # 注释说明：以下用于监督学习任务
        # 构建原始标签文件的完整路径（用于测试）
        self.labels_path = os.path.join(data_root, 'labels.h5')
        # 构建视频级标签文件的完整路径（用于多实例学习）
        self.dir_labels_path = os.path.join(data_root, 'mil_labels.h5')
        # 构建背景噪声标签文件的完整路径（仅包含背景）
        self.dir_labels_bg_path = os.path.join(data_root, 'labels_noisy.h5')
        # 构建样本顺序文件的完整路径，文件名根据分割类型动态生成
        self.sample_order_path = os.path.join(data_root, f'{split}_order.h5')
        # 标记HDF5文件是否已打开，初始状态为未打开
        self.h5_isOpen = False

    def __getitem__(self, index):
        """
        获取指定索引的数据样本
        
        参数:
        index: 样本索引，整数类型
        
        返回:
        visual_feat: 视觉特征，numpy数组或tensor
        audio_feat: 音频特征，numpy数组或tensor  
        label: 样本标签，numpy数组或tensor
        """
        
        # 如果HDF5文件尚未打开，则打开所有需要的HDF5文件
        if not self.h5_isOpen:
            # 打开样本顺序HDF5文件并获取顺序数组
            self.sample_order = h5py.File(self.sample_order_path, 'r')['order']
            # 打开干净视觉特征HDF5文件并获取数据集
            self.visual_feature = h5py.File(self.visual_feature_path, 'r')['avadataset']
            # 打开干净音频特征HDF5文件并获取数据集
            self.audio_feature = h5py.File(self.audio_feature_path, 'r')['avadataset']
            # 打开视频级标签HDF5文件并获取数据集
            self.clean_labels = h5py.File(self.dir_labels_path, 'r')['avadataset']

            # 如果是训练阶段，加载噪声数据
            if self.split == 'train':
                # 打开背景噪声标签HDF5文件并获取数据集
                self.negative_labels = h5py.File(self.dir_labels_bg_path, 'r')['avadataset']
                # 打开噪声视觉特征HDF5文件并获取数据集
                self.negative_visual_feature = h5py.File(self.noisy_visual_feature_path, 'r')['avadataset']
                # 打开噪声音频特征HDF5文件并获取数据集
                self.negative_audio_feature = h5py.File(self.noisy_audio_feature_path, 'r')['avadataset']

            # 如果是测试阶段，加载原始标签
            if self.split == 'test':
                # 打开原始标签HDF5文件并获取数据集
                self.labels = h5py.File(self.labels_path, 'r')['avadataset']

            # 标记HDF5文件已打开，避免重复打开
            self.h5_isOpen = True

        # 获取干净样本的数量
        clean_length = len(self.sample_order)
        
        # 如果索引大于等于干净样本数量，说明要获取噪声样本
        if index >= clean_length:
            # 计算噪声样本的有效索引
            valid_index = index - clean_length
            # 从噪声视觉特征数据集中获取对应样本的视觉特征
            visual_feat = self.negative_visual_feature[valid_index]
            # 从噪声音频特征数据集中获取对应样本的音频特征
            audio_feat = self.negative_audio_feature[valid_index]
            # 从噪声标签数据集中获取对应样本的标签
            label = self.negative_labels[valid_index]
        else:
            # 处理干净样本或测试阶段的样本
            # 根据索引从样本顺序数组中获取实际的样本索引
            sample_index = self.sample_order[index]
            # 从干净视觉特征数据集中获取对应样本的视觉特征
            visual_feat = self.visual_feature[sample_index]
            # 从干净音频特征数据集中获取对应样本的音频特征
            audio_feat = self.audio_feature[sample_index]
            
            # 根据数据集分割类型选择标签来源
            if self.split == 'train':
                # 训练阶段使用视频级标签
                label = self.clean_labels[sample_index]
            else:
                # 测试阶段使用原始标签
                label = self.labels[sample_index]

        # 返回样本的视觉特征、音频特征和标签
        return visual_feat, audio_feat, label

    def __len__(self):
        """
        获取数据集的样本总数量
        
        返回:
        length: 数据集的样本总数，整数类型
        """
        
        # 根据数据集分割类型计算不同的样本数量
        if self.split == 'train':
            # 训练阶段：干净样本数量 + 噪声样本数量
            # 打开样本顺序HDF5文件并获取顺序数组
            sample_order = h5py.File(self.sample_order_path, 'r')['order']
            # 打开噪声标签HDF5文件并获取数据集
            noisy_labels = h5py.File(self.dir_labels_bg_path, 'r')['avadataset']
            # 计算总样本数量
            length = len(sample_order) + len(noisy_labels)
        elif self.split == 'test':
            # 测试阶段：仅干净样本数量
            # 打开样本顺序HDF5文件并获取顺序数组
            sample_order = h5py.File(self.sample_order_path, 'r')['order']
            # 获取样本数量
            length = len(sample_order)
        else:
            # 不支持的其他分割类型，抛出未实现错误
            raise NotImplementedError

        # 返回样本总数量
        return length